#!/usr/bin/env python3

import cv2
import rospy
from sensor_msgs.msg import Image
from geometry_msgs.msg import Vector3Stamped
from cv_bridge import CvBridge, CvBridgeError
import os
import time


class ImageSaver:
    def __init__(
        self,
        topic_name="/camera/color/image_raw",
        state_topic="/servo/state",
        save_interval=0.1,
        output_dir="./images",
    ):
        self.bridge = CvBridge()
        self.last_save_time = 0
        self.save_interval = save_interval
        self.output_dir = output_dir
        os.makedirs(self.output_dir, exist_ok=True)
        self.image_count = 0
        self.angular_velocity = 0.0  # Initialize with non-zero value

        rospy.init_node("image_saver_node", anonymous=True)

        # Subscribe to image topic
        rospy.Subscriber(topic_name, Image, self.image_callback)
        rospy.loginfo(f"Subscribed to image topic: {topic_name}")

        # Subscribe to state topic
        rospy.Subscriber(state_topic, Vector3Stamped, self.state_callback)
        rospy.loginfo(f"Subscribed to state topic: {state_topic}")

    def image_callback(self, msg):
        # Check if angular velocity y is zero
        if self.angular_velocity != 0.0:
            return

        try:
            # Convert ROS Image message to OpenCV image
            cv_image = self.bridge.imgmsg_to_cv2(msg, "bgr8")
        except CvBridgeError as e:
            rospy.logerr(f"CvBridge Error: {e}")
            return

        # Get current time in seconds
        current_time = time.time()

        # Check if enough time has passed since the last image was saved
        if current_time - self.last_save_time >= self.save_interval:
            image_path = os.path.join(
                self.output_dir, f"frame_{self.image_count:04d}.jpg"
            )
            cv2.imwrite(image_path, cv_image)
            rospy.loginfo(f"Saved image: {image_path}")
            self.image_count += 1
            self.last_save_time = current_time

    def state_callback(self, msg):
        # Update angular velocity y value
        self.angular_velocity = msg.vector.y

    def run(self):
        rospy.spin()


if __name__ == "__main__":
    # Save in relative path
    dir = os.path.join(os.path.dirname(__file__), "../images")
    image_saver = ImageSaver(
        topic_name="/camera/color/image_raw",
        state_topic="/servo/state",
        save_interval=0.1,
        output_dir=dir,
    )
    image_saver.run()
